feat(rocm): V1 sampling ops for Hermes EngineCore - #234

TESTERS WANTED — BF16 / FP16 Gemma-4 MoE on ROCm

Lab proof for this PR was FP8 Gemma-4-26B-A4B MoE on 2× Radeon AI PRO R9700 (gfx1201). The code path here is dtype-agnostic (V1 sampler ops + Hermes max_tokens=-1), but we cannot reliably soak BF16 MoE on this box:

Please test if you have

• Gemma-4-26B MoE (or similar MoE) weights in BF16 or FP16 (not only FP8)
• ROCm multi-GPU with enough VRAM (e.g. 2×48, 1×80+, or comfortable 2×32 with short ctx)
• Enough host RAM to load without thrashing

Minimal matrix

1. Short chat: Paris / READY
2. Hermes-like: max_tokens=-1, temperature/top_p > 0
3. Title-gen style: allowed_token_ids present
4. Medium stream (tools optional) — confirm no no kernel for op … on device type 5 and no stopped AsyncLLM after first failure

Report back

• GPU(s) + ROCm version
• Model id / dtype
• max_model_len / num_blocks
• Pass/fail per step + any log lines with engine-fatal / no kernel for op

Related: serve KV fail-fast #227; FP8 MoE prefill/resident package #228 (FP8-proven; MoE structure notes there).

Thanks — BF16/FP16 MoE ROCm coverage has to come from the community; FP8 is what we can gate on RDNA4 consumer hardware day-to-day.

Addressing CHANGES_REQUESTED (localai-org-maint-bot)

Automated sampler coverage

Added CPU↔ROCm parity cases in tests/vt/test_ops_sample.cpp (skip cleanly when ROCm is not linked):

Lab run on this box (HIP build, dual R9700):

./tests/test_ops_sample --test-case='*ROCm*'
test cases: 2 passed | assertions: 4108 passed | SUCCESS

CPU CI still exercises the existing hand-oracle sampler suite; ROCm block is opt-in like the CUDA parity block.

Docs gates

• docs/USAGE.md: max_tokens=-1 → VT_SERVER_MAX_NEW_TOKENS, ROCm V1 sampler surface note
• Local: check-doc-checkpoint OK, check-agent-record OK

Merge order (unchanged)

Land #234 first → rebase #227 / #228.

Happy to add CUDA-style large-vocab ROCm greedy case or penalty/logit-bias matrix next if you want a wider surface before merge.

Fixed agent-record: FEATURES.md Gemma4 ROCm cell was 258 chars (limit 220). Shortened; check-public-doc-tables OK.

Verified on CUDA and it's clean there — one substantive change requested before I merge.

Verified for you (GB10, sm_121a, rebased on current main)

test_ops_sample            279749/279749
test_openai_serving             534/534
test_qwen27_paged_engine        235/235   (SACRED)
test_qwen36_paged_engine        315/315   (SACRED)

So the ROCm sampling ops are inert on CUDA and the shared OpenAI-layer change doesn't disturb the gate models. Nice, self-contained work.

Change requested: max_tokens < 1 should map to UNSET, not to 4096

protocol.cpp:570-575:

// Hermes / some clients send max_tokens=-1 meaning "default/unlimited".
// SamplingParams::PostInit rejects <1; map to a lab-safe default (4096).
if (req_max.has_value() && *req_max < 1) req_max = 4096;

The diagnosis is right — Hermes really does send -1 for "unlimited" and PostInit really does reject <1. But 4096 is the wrong target, because we already implement "unlimited" correctly, one layer down:

src/vllm/v1/engine/input_processor.cpp:175-180

// params is already our clone (passed by value). If unset max_tokens, then
// generate up to the max_model_len (input_processor.py:317-321).
if (!params.max_tokens.has_value()) {
  const int64_t seq_len = static_cast<int64_t>(prompt_token_ids.size());
  params.max_tokens = static_cast<int>(max_model_len_ - seq_len);
}

That is a 1:1 mirror of vLLM (input_processor.py:317-321, cited in our own comment), and the very next line in your file already routes an absent value there: sp.max_tokens = req_max.has_value() ? req_max : default_max_tokens;.

So the one-character-cheaper fix says what the client meant:

// -1 means "no client-side limit", which is exactly "unset": the engine then
// generates to max_model_len - seq_len (input_processor.cpp:175-180, mirroring
// vllm input_processor.py:317-321).
if (req_max.has_value() && *req_max < 1) req_max = std::nullopt;

Why it matters rather than being a style nit: with max_model_len 32768 and a 2k prompt, a Hermes client asking for "unlimited" currently gets silently truncated at 4096 with finish_reason=length and no way to tell it apart from a real limit. It also hardcodes a policy number into the protocol layer where VT_SERVER_MAX_NEW_TOKENS already exists as the tunable, and it diverges from vLLM, which we mirror here.

Also, lines 573-575 are dead code — an if whose body is only a comment:

if (req_max.has_value() && *req_max > 4096) {
  // soft cap; serving_chat may clamp further via VT_SERVER_MAX_NEW_TOKENS
}

It compiles to nothing and reads like a cap that isn't there. Either drop it or make it real.

Make those two edits and I'll re-run the same set and merge — everything else here is good to go.

Heads-up: the max_tokens clamp is now fixed on main as 2d2bdd04 (row/SERVE-MAXTOKENS-UNSET), so please drop it from this PR and rebase — otherwise the three PRs conflict over one spot.

A non-positive max_tokens (or max_completion_tokens) is now treated as unset on both request paths, which routes to max_model_len - seq_len in input_processor.cpp:175-180, mirroring vllm input_processor.py:317-321. No constant, no truncation of the long-context request that asked to be left unlimited.

RED before / GREEN after on dgx, the new test built against unpatched main:

before: THREW 'max_tokens must be at least 1, got -1';  28/29 cases, 2 assertions failed
after:  29/29 cases, 179/179 assertions

test_openai_protocol now pins the semantics (−1 and 0 unset on both paths, positive unchanged, and −1 yielding to a serving default). Semantics documented in docs/USAGE.md.

Everything else in your PR stands — this is only about the shared hunk.

Rebase + drop max_tokens clamp (addressing CHANGES_REQUESTED / main drift)

Rebuilt tip-of-main (a0fa12c7) as single commit 42ba78e2.

Dropped (now on main)

• max_tokens=-1 → 4096 clamp in protocol.cpp / serving_chat.cpp — main already treats non-positive as unset (2d2bdd04 / SERVE-MAXTOKENS-UNSET) and fills max_model_len - seq_len in input_processor. Keeping our clamp would fight that and conflict with fix(server): fail-fast unschedulable KV waits #227/feat(gemma4/rocm): FP8 resident MoE + SharedK-WMMA prefill (RDNA4) #228.

Kept (this PR only)

• Full ROCm V1 sampler surface: rocm_sample.hip + allowed/token-mask helpers in rocm_dense_basic.hip + RegisterOp in rocm_ops.hip
• CMakeLists: add rocm_sample.hip only (no HIP -O / Triton comment thrash)
• CPU↔ROCm parity cases in tests/vt/test_ops_sample.cpp (skip when ROCm unlinked)
• USAGE note + FEATURES Gemma4 ROCm cell shortened under public-table limit

Local gates (on this tip)

check-device-leakage, check-public-doc-tables, check-doc-checkpoint, check-env-doc, check-agent-record, check-pr-size — OK.

Please re-run CI on the force-pushed head. Happy to split further if review still wants sampler kernels alone vs docs.

Rebased onto latest main (f86e4cc0) — tip 8ea62b21. Local gates green (doc-checkpoint, env-doc, device-leakage, pr-size). Please re-run CI.

Merged. This closes a real bug: ROCm registered no sampling ops at all, so EngineCore fatally hit "no kernel for op" the moment prefill finished on AMD. The full V1 surface is now registered for DeviceType::kROCM — temperature, top-k/top-p, probs, logprobs, random sample, penalties, min-p, logit bias, allowed-token-ids and the sparse bad-words mask.

Your red agent-record and pr-size checks were the same fork artifact as #283, not a defect: both refused with base must be an ancestor of head because the branch was never rebased, so neither could compute a commit range. sanitize-cpu (address,undefined) is the pre-existing main baseline present on every open PR. Nothing red was attributable to this change.

Since no CI job compiles ROCm — there are no AMD runners — the green checks are orthogonal either way, so I verified the port mechanically instead of by reading it. I normalized src/vt/rocm/rocm_sample.hip against src/vt/cuda/cuda_sample.cu (folding away namespace, hip/cuda prefixes, kernel-name suffixes and whitespace) and diffed. Every remaining difference is cosmetic: entry-point names, error-string prefixes, line wrapping, and CUDA's trailing explanatory comments.

The numerics are identical, including every piece that decides token identity:

• SplitMix64 constants
• ExpNoise's (r >> 11) + 1 over 9007199254740993.0 mantissa construction
• the gumbel/exp-noise argmax
• the temperature !all_random && t < kSamplingEps guard
• flashinfer's two-pivot sort-free bracket search, same kThreshMaxIter = 64, same min_gt_low / max_le_high snapping

That's a genuine 1:1 port rather than a reimplementation, which is what this project's porting rule asks for and what makes it reviewable without an AMD box.

One thing I reviewed and deliberately accepted as-is: RandomSampleK launches <<<n, 1>>> — one thread per row scanning the whole vocab serially. That's not a defect here, because cuda_sample.cu launches it exactly the same way, so changing it would be a divergence from the reference. It is a real optimization target for whoever takes ROCm sampling past correctness-grade, and worth its own issue when someone does.

Tests skip cleanly via HasRocm() when no ROCm backend is registered, so they're inert on CPU CI and assert real CPU-vs-ROCm parity on your gfx1201 pair. I also checked that the docs/USAGE.md non-positive max_tokens claim is backed by code already on main (protocol.cpp:525) rather than asserted by this PR without landing it.

Separately — I've left detailed review notes on #228 asking for that one to be split rather than merged. Different situation entirely; this one was clean.

Tip-of-main rebuild (back on track)

Rebuilt as single commit on current main (3286f6fd) → tip a418791e.

• Payload only: rocm_sample.hip + ops registration + dense_basic hooks + CPU↔ROCm parity tests + minimal FEATURES/USAGE note
• No max_tokens clamp (already on main via 2d2bdd04)
• Docs restored from main (no FEATURES/USAGE clobber)
• Local gates green: device-leakage, doc-checkpoint, public-doc-tables, env-doc, readme, commit-trailers, pr-size, agent-record

Please re-run CI. Prior red agent-record/pr-size were base-not-ancestor.

CI hygiene 2026-08-11: rebased onto origin/main (60e71a0). Local gates green (doc-checkpoint, env-doc, agent-record, pr-size, device-leakage). Land order remains #234 → #227 → #316 → #317.

Rebased onto origin/main c70f42b9 (main moved +2 during CI). Local gates green. Land order unchanged: #234 → #227 → #316 → #317. PR bodies refreshed (RandomSample note on #234; #228 closed in land path).

@localai-org-maint-bot Status check against CHANGES_REQUESTED / localai-bot notes:

1. CPU/ROCm parity tests — tests/vt/test_ops_sample.cpp adds HasRocm() cases covering greedy/temperature/top-k-p/allowed_ids/logprobs and min_p/penalties (skip cleanly when no ROCm backend — inert on CPU CI).
2. max_tokens→4096 — removed from this PR; main owns non-positive → unset (2d2bdd04). No protocol.cpp hunk on the tip.
3. PR is still OPEN (head b067cd3a) — treating the earlier “Merged.” comment as review prose, not GitHub merge state.

Sanitize ASan/TSan remain ambient-main class. Product jobs (build-test-cpu, cuda-fat, vulkan) were green on the last full run.

Landed on main as 64366dc09a9b — the full V1 sampler registration on ROCm is on main and verified present.

Closing manually: it landed as a rebased copy of this branch, so the head SHA here never became reachable from main and GitHub could not auto-close it. Nothing is outstanding; the PR was just left showing OPEN.

Thanks for the contribution.